from models.roberta_model_multiclass import BertMultiClass
from data_utils.bert_multi_class_data_roberta import get_train_val_data_loader_cross, get_test_loader_by_split_id, \
    get_test_loader
import torch
from config import conf
import pandas as pd
from os.path import join, exists
import os
FEATURE_ROOT_DIR = conf.get('linux_dir', 'feature_root_dir')
DEVICE_ID = 1#conf.get('gpu', 'device_id')



def train(i):
    device = torch.device("cuda:%s" % (DEVICE_ID) if torch.cuda.is_available() else "cpu")

    train_dataloader, val_dataloader, tokenizer = get_train_val_data_loader_cross(device, 4, test_number=i,
                                                                                  cross_number=9, shuffle=True,
                                                                                  maxlen=400)
    bert_attn_model = BertMultiClass(learning_rate=5e-6, model_name='multi_class_cross_roberta').to(device)
    bert_attn_model.train_epochs(train_dataloader, val_dataloader, num_epoch=15, early_stop=5)
    print('test ', i)
    bert_attn_model.load_best_model()
    test_loader = get_test_loader(device, 4, maxlen=400)
    results = list(bert_attn_model.inference(test_loader))
    save_path = join(bert_attn_model.model_save_dir, 'data', 'raw_rs.csv')
    pd.DataFrame(results, columns=['id', 'predict_labels']).to_csv(save_path, index=False)
    print('saved ', save_path)


def predict_for_new_features(i):
    device = torch.device("cuda:%s" % (DEVICE_ID) if torch.cuda.is_available() else "cpu")
    test_loader = get_test_loader_by_split_id(device, 16, test_number=i)
    bert_attn_model = BertMultiClass(learning_rate=5e-6, model_name='multi_class_cross_roberta', version_id=i).to(
        device)
    bert_attn_model.load_best_model()
    results = bert_attn_model.inference_for_features(test_loader)
    save_path = join(FEATURE_ROOT_DIR, bert_attn_model.model_name, 'feature_split %d' % (i))
    save_dir = os.path.dirname(save_path)
    if not exists(save_dir):
        os.makedirs(save_dir)
    pd.DataFrame(results, columns=['id', 'predict_features']).to_csv(save_path, index=False)
    print('saved at %s' % save_path)

def predict_for_new_test_features(i):
    device = torch.device("cuda:%s" % (DEVICE_ID) if torch.cuda.is_available() else "cpu")
    test_loader = get_test_loader(device, 16)
    bert_attn_model = BertMultiClass(learning_rate=5e-6, model_name='multi_class_cross_roberta', version_id=i).to(
        device)
    bert_attn_model.load_best_model()
    results = bert_attn_model.inference_for_features(test_loader)
    save_path = join(FEATURE_ROOT_DIR, bert_attn_model.model_name, 'test_features_round2_version %d' % (i))
    save_dir = os.path.dirname(save_path)
    if not exists(save_dir):
        os.makedirs(save_dir)
    pd.DataFrame(results, columns=['id', 'predict_test_features']).to_csv(save_path, index=False)
    print('saved at %s' % save_path)


if __name__ == '__main__':
    for i in range(4, 10):
        train(i)
        torch.cuda.empty_cache()
    for i in range(1, 10):
        predict_for_new_features(i)
        torch.cuda.empty_cache()
    for i in range(1, 10):
        predict_for_new_test_features(i)
        torch.cuda.empty_cache()
